{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import nuscenes\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nusc = nuscenes.NuScenes(version='v1.0-trainval', dataroot='/data/Datasets/nuScenes', verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "sample = nusc.sample[13]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nusc.render_sample_data(nusc.get('sample_data', sample['data']['LIDAR_TOP'])['token'], nsweeps=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "lidar_path, boxes, _ = nusc.get_sample_data(nusc.get('sample_data', sample['data']['LIDAR_TOP'])['token'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lidar_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "points = np.fromfile(lidar_path, dtype=np.float32).reshape([-1, 5])[:, :4].T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "points.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from nuscenes.utils.geometry_utils import transform_matrix\n",
    "from pyquaternion import Quaternion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "curr_sd_rec = nusc.get('sample_data', nusc.get('sample', sample['token'])['data']['LIDAR_TOP'])\n",
    "ego_pose = nusc.get('ego_pose', curr_sd_rec['ego_pose_token'])\n",
    "calibrated_sensor = nusc.get('calibrated_sensor', curr_sd_rec['calibrated_sensor_token'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "global_from_car = transform_matrix(\n",
    "                    ego_pose['translation'],\n",
    "                    Quaternion(ego_pose['rotation']),\n",
    "                    inverse=False)\n",
    "car_from_current = transform_matrix(\n",
    "                    calibrated_sensor['translation'],\n",
    "                    Quaternion(calibrated_sensor['rotation']),\n",
    "                    inverse=False)\n",
    "\n",
    "tm = np.dot(global_from_car, car_from_current)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nbr_points = points.shape[1]\n",
    "points[:3, :] = tm.dot(np.vstack((points[:3, :], np.ones(nbr_points))))[:3, :]\n",
    "points = points.T\n",
    "print(points.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(points.shape[1]):\n",
    "    print(points[:, i].min(), points[:, i].max())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.scatter(points[:, 0],points[:, 1])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "scene = nusc.get('scene', nusc.get('sample', sample['token'])['scene_token'])\n",
    "scene_token = scene['token']\n",
    "scene_record = nusc.get('scene', scene_token)\n",
    "log_record = nusc.get('log', scene_record['log_token'])\n",
    "map_record = nusc.get('map', log_record['map_token'])\n",
    "map_mask = map_record['mask']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "cropped = map_mask.mask()[8940:10695, 5599: 6969]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(cropped)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ego_pose['translation']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "map_pose = np.concatenate(map_mask.to_pixel_coords(ego_pose['translation'][0], ego_pose['translation'][1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_pose.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_pose"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "filtered = map_mask.is_on_mask(x=points[:, 0], y=points[:, 1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "filtered_points = points[filtered]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filtered_points.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.scatter(filtered_points[:, 0],filtered_points[:, 1])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "x_min = map_pose[0] - int(50.4/0.1)\n",
    "x_max = map_pose[0] + int(50.4/0.1)\n",
    "y_min = map_pose[1] - int(50.4/0.1)\n",
    "y_max = map_pose[1] + int(50.4/0.1)\n",
    "cropped2 = map_mask.mask()[y_min:y_max, x_min:x_max, ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(cropped2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
